use tonic::{Request, Response, Status, Streaming};
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::RwLock;
use std::time::{Duration};
use tonic::codegen::tokio_stream::{Stream, StreamExt};
use tracing::{info, error, debug};

pub mod echo {
    tonic::include_proto!("echo");
}

use echo::{
    echo_service_server::{EchoService, EchoServiceServer},
    EchoRequest, EchoResponse, HealthCheckRequest, HealthCheckResponse,
};

#[derive(Debug)]
pub struct ConnectionStats {
    pub total_connections: usize,
    pub broken_pipes: usize,
    pub successful_reconnects: usize,
    pub active_connections: usize,
}

impl ConnectionStats {
    pub fn new() -> Self {
        Self {
            total_connections: 0,
            broken_pipes: 0,
            successful_reconnects: 0,
            active_connections: 0,
        }
    }
}

#[derive(Debug)]
pub struct MyEchoService {
    stats: Arc<RwLock<ConnectionStats>>,
    server_id: String,
}

impl MyEchoService {
    pub fn new() -> Self {
        Self {
            stats: Arc::new(RwLock::new(ConnectionStats::new())),
            server_id: format!("server-{}", rand::random::<u16>()),
        }
    }
}

#[tonic::async_trait]
impl EchoService for MyEchoService {
    async fn echo(&self, request: Request<EchoRequest>) -> Result<Response<EchoResponse>, Status> {
        let req = request.into_inner();

        info!("收到Echo请求: {}, 序列: {}", req.message, req.sequence);

        let response = EchoResponse {
            message: format!("Echo: {}", req.message),
            sequence: req.sequence,
            server_id: self.server_id.clone(),
            timestamp: chrono::Utc::now().timestamp_millis(),
        };

        {
            let mut stats = self.stats.write().await;
            stats.total_connections += 1;
            stats.active_connections += 1;
        }

        // 模拟处理延迟
        tokio::time::sleep(Duration::from_millis(50)).await;

        info!("发送Echo响应: {}, 序列: {}", response.message, response.sequence);

        Ok(Response::new(response))
    }
    type EchoStreamStream = Pin<Box<dyn Stream<Item=Result<EchoResponse, Status>> + Send + 'static>>;

    async fn echo_stream(&self, request: Request<EchoRequest>) -> Result<Response<Self::EchoStreamStream>, Status> {
        let req = request.into_inner();
        let server_id = self.server_id.clone(); // 提前克隆到局部变量

        info!("开始流式Echo: {}, 序列: {}", req.message, req.sequence);

        let stream = async_stream::stream! {
            for i in 0..5 {
                let response = EchoResponse {
                    message: format!("流式Echo {}: {}", i + 1, req.message),
                    sequence: req.sequence * 10 + i as i32,
                    server_id: server_id.clone(), // 使用局部变量而非self引用
                    timestamp: chrono::Utc::now().timestamp_millis(),
                };

                yield Ok(response);

                tokio::time::sleep(Duration::from_millis(200)).await;
            }
        };
        Ok(Response::new(Box::pin(stream)))
    }

    async fn echo_client_stream(&self, request: Request<Streaming<EchoRequest>>) -> Result<Response<EchoResponse>, Status> {
        let mut stream = request.into_inner();
        let mut message_count = 0;
        let mut combined_message = String::new();

        while let Some(req) = stream.message().await? {
            message_count += 1;
            combined_message.push_str(&format!("[{}] ", req.message));
        }

        let response = EchoResponse {
            message: format!("收到 {} 条消息: {}", message_count, combined_message),
            sequence: -1,
            server_id: self.server_id.clone(),
            timestamp: chrono::Utc::now().timestamp_millis(),
        };

        info!("客户端流处理完成，总共 {} 条消息", message_count);

        Ok(Response::new(response))
    }

    type EchoBidirectionalStream = Pin<Box<dyn Stream<Item=Result<EchoResponse, Status>> + Send + 'static>>;

    async fn echo_bidirectional(&self, request: Request<Streaming<EchoRequest>>) -> Result<Response<Self::EchoBidirectionalStream>, Status> {
        let mut stream = request.into_inner();
        let server_id = self.server_id.clone();  // 提前克隆共享数据

        let output_stream = async_stream::stream! {
            let mut sequence = 0;

            while let Some(result) = stream.next().await {
                match result {
                    Ok(req) => {
                        let response = EchoResponse {
                            message: format!("双向Echo: {}", req.message),
                            sequence,
                            server_id: server_id.clone(),  // 使用已克隆的副本
                            timestamp: chrono::Utc::now().timestamp_millis(),
                        };
                        sequence += 1;
                        yield Ok(response);
                    }
                    Err(e) => {
                        error!("处理双向流时出错: {}", e);
                        // 直接返回错误状态，而不是创建新的错误
                        yield Err(Status::internal(format!("处理流时出错: {}", e)));
                    }
                }
            }
        };

        Ok(Response::new(Box::pin(output_stream)))
    }

    async fn health_check(&self, request: Request<HealthCheckRequest>) -> Result<Response<HealthCheckResponse>, Status> {
        let req = request.into_inner();

        debug!("健康检查请求: {}", req.service);

        let response = HealthCheckResponse {
            status: HealthCheckResponse::status(&Default::default()).into(),
            server_info: format!("服务器ID: {}, 活跃连接: {}", self.server_id, self.stats.read().await.active_connections),
        };

        Ok(Response::new(response))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt::init();

    let addr = "[::1]:50051".parse()?;
    let echo_service = MyEchoService::new();
    let stats = Arc::clone(&echo_service.stats);

    info!("gRPC服务器启动，监听地址: {}", addr);
    info!("服务器ID: {}", echo_service.server_id);

    // 启动统计信息报告
    let stats_clone = Arc::clone(&stats);
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(Duration::from_secs(10)).await;

            let stats = stats_clone.read().await;
            info!("服务器统计 - 总连接: {}, Broken pipes: {}, 成功重连: {}, 活跃连接: {}", stats.total_connections,
                stats.broken_pipes,
                stats.successful_reconnects,
                stats.active_connections
            );
        }
    });

    tonic::transport::Server::builder()
        .add_service(EchoServiceServer::new(echo_service))
        .serve(addr)
        .await?;

    Ok(())
}
